import torch.nn as nn
import torch.nn.functional as F


class CIFAR10Model(nn.Module):
    def __init__(self):
        super(CIFAR10Model, self).__init__()
        # 卷积层1: 输入通道3(彩色图像), 输出通道32, 3x3卷积核
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        # 池化层1: 2x2最大池化
        self.pool1 = nn.MaxPool2d(2, 2)

        # 卷积层2: 输入通道32, 输出通道64, 3x3卷积核
        self.conv2 = nn.Conv2d(32, 640, kernel_size=3, padding=1)
        # 池化层2: 2x2最大池化
        self.pool2 = nn.MaxPool2d(2, 2)

        # 卷积层3: 输入通道64, 输出通道64, 3x3卷积核
        self.conv3 = nn.Conv2d(640, 1280, kernel_size=3, padding=1)
        # 池化层2: 2x2最大池化
        self.pool3 = nn.MaxPool2d(2, 2)

        # 卷积层2: 输入通道32, 输出通道64, 3x3卷积核
        self.conv4 = nn.Conv2d(1280, 2560, kernel_size=3, padding=1)
        # 卷积层3: 输入通道64, 输出通道64, 3x3卷积核
        self.conv5 = nn.Conv2d(2560, 1280, kernel_size=3, padding=1)

        # 全连接层
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(1280 * 4 * 4, 64)  # 注意计算输入维度
        self.fc2 = nn.Linear(64, 10)  # 输出10个类别

    def forward(self, x):
        # 卷积层1 + ReLU + 池化1
        x = F.relu(self.conv1(x))  # [b,3,wide,high] -> [b,32,wide,high]
        x = self.pool1(x)  # [b,32,wide,high] -> [b,32,wide/2,high/2]

        # 卷积层2 + ReLU + 池化2
        x = F.relu(self.conv2(x))  # [b,32,wide/2,high/2] -> [b,64,wide/2,high/2]
        x = self.pool2(x)  # [b,32,wide/2,high/2] -> [b,64,wide/4,high/4]

        # 卷积层3 + ReLU
        x = F.relu(self.conv3(x))  # [b,64,wide/2,high/2] -> [b,64,wide/4,high/4]
        x = self.pool3(x)  # [b,32,wide/2,high/2] -> [b,64,wide/4,high/4]

        x = F.relu(self.conv4(x))  # [b,64,wide/2,high/2] -> [b,64,wide/4,high/4]
        x = F.relu(self.conv5(x))  # [b,64,wide/2,high/2] -> [b,64,wide/4,high/4]

        # 展平 + 全连接层
        x = self.flatten(x)  # [b,64,wide/4,high/4] -> [b,64*wide/4*high/4]
        x = F.relu(self.fc1(x))  # [b,64*wide/4*high/4] -> [b,64]
        x = self.fc2(x)  # [b,64] -> [b,10]
        return x
